package com.ybear.ybutils.utils;


import android.util.Base64;

import androidx.annotation.NonNull;
import androidx.annotation.Nullable;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.nio.charset.Charset;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * AES 加密解密工具
 */
public class AESUtil {
    /**
     加密模式
     */
    @Retention( RetentionPolicy.SOURCE )
    public @interface Mode {
        String ECB = "ECB";
        String CBC = "CBC";
        String CTR = "CTR";
        String OFB = "OFB";
        String CFB = "CFB";
    }

    /**
     填充
     */
    @Retention( RetentionPolicy.SOURCE )
    public @interface Padding {
        String PKCS5_PADDING = "PKCS5Padding ";
        String PKCS7_PADDING = "PKCS7Padding ";
        String ZERO_PADDING = "ZeroPadding";
        String ISO_10126 = "ISO10126";
        String ANSIX_923 = "ANSIX923";
        String NO_PADDING = "NoPadding";
    }

    private static String mMode = Mode.ECB;
    private static String mPadding = Padding.NO_PADDING;
    private static Charset mDefaultCharset = Charset.defaultCharset();
    private static int mDefaultBase64Flags = Base64.DEFAULT;
    
    /**
     加密模式
     @param mode    {@link Mode}
     */
    public static void setMode(@Mode String mode) { mMode = mode; }

    /**
     填充
     @param padding {@link Padding}
     */
    public static void setPadding(@Padding String padding) { mPadding = padding; }

    /**
     字符集
     @param charset GB-2312, UTF-8...
     */
    public static void setCharset(Charset charset) { mDefaultCharset = charset; }

    /**
     Base64 flag
     @param flags   controls certain features of the encoded output.
      *               Passing {@code DEFAULT} results in output that
      *               adheres to RFC 2045.
     */
    public static void setBase64Flags(int flags) {
        mDefaultBase64Flags = flags;
    }

    /**
     切换为Java环境（单元测试的时候Android部分代码无法使用）
     @param mode true:Java环境， false:Android环境
     */
    public static void setJavaMode(boolean mode) { Base64Utils.setJavaMode( mode ); }

    /**
     创建用于AES加密解密的Cipher
     @param key         密钥
     @param iv          偏移量
     @param opMode      加密/解密
     @return            {@link Cipher}
     */
    @Nullable
    private static Cipher createCipher(String key, String iv, int opMode) {
        Cipher cipher;
        String type = "AES";
        try {
            //加密模式
            cipher = Cipher.getInstance( type + "/" + mMode + "/" + mPadding );
        } catch(NoSuchAlgorithmException | NoSuchPaddingException e) {
            e.printStackTrace();
            return null;
        }
        //密钥
        SecretKeySpec skSpec = new SecretKeySpec( key.getBytes( mDefaultCharset ), type );
        //偏移量（ECB模式不支持iv）
        IvParameterSpec ivpSpec = Mode.ECB.equals( mMode ) ?
                null :
                new IvParameterSpec( iv.getBytes( mDefaultCharset ) );
        if( ivpSpec == null ) {
            try {
                cipher.init( opMode, skSpec );
            } catch(InvalidKeyException e) {
                e.printStackTrace();
            }
        }else {
            try {
                cipher.init( opMode, skSpec, ivpSpec );
            } catch(InvalidAlgorithmParameterException | InvalidKeyException e) {
                e.printStackTrace();
            }
        }
        return cipher;
    }

    /**
     加密方法
     @param data      数据
     @param key       密钥
     @param iv        偏移量
     @return          结果
     */
    @NonNull
    public static String encrypt(String data, String key, String iv) {
        try {
            Cipher cipher = createCipher( key, iv, Cipher.ENCRYPT_MODE );
            if( cipher == null ) return "";
            int blockSize = cipher.getBlockSize();
            byte[] dataBytes = data.getBytes( mDefaultCharset );
            int plaintextLength = dataBytes.length;

            if( plaintextLength % blockSize != 0 ) {
                plaintextLength = plaintextLength + ( blockSize - ( plaintextLength % blockSize ) );
            }
            byte[] plaintext = new byte[ plaintextLength ];
            System.arraycopy( dataBytes, 0, plaintext, 0, dataBytes.length );
            return Base64Utils.encodeBase64(
                    cipher.doFinal( plaintext ), mDefaultCharset, mDefaultBase64Flags
            );
        } catch (BadPaddingException | IllegalBlockSizeException | AssertionError e) {
            e.printStackTrace();
        }
        return "";
    }
    /**
     加密方法
     @param data      数据
     @param key       密钥
     @return          结果
     */
    public static String encrypt(String data, String key) { return encrypt( data, key, key ); }

    /**
     解密方法
     @param data 数据
     @param key  密钥
     @param iv   偏移量
     @return 解密的结果
     */
    @NonNull
    public static String decrypt(String data, String key, String iv) {
        try {
            Cipher cipher = createCipher( key, iv, Cipher.DECRYPT_MODE );
            return new String(
                    cipher == null ?
                            new byte[ 0 ] :
                            cipher.doFinal(Base64Utils.decodeBase64(
                                    data, mDefaultCharset, mDefaultBase64Flags
                            )),
                    mDefaultCharset
            ).trim();
        } catch (BadPaddingException | IllegalBlockSizeException | IllegalArgumentException e) {
            e.printStackTrace();
        }
        return "";
    }

    /**
     解密方法
     @param data 数据
     @param key  密钥
     @return 解密的结果
     */
    public static String decrypt(String data, String key) { return decrypt( data, key, key ); }
}